from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
from insertByMilvus import insert_texts_to_milvus


def extract_text_from_pdf(filename, page_numbers=None, min_line_length=1):
    '''从 PDF 文件中（按指定页码）提取文字'''
    paragraphs = []
    buffer = ''
    full_text = ''
    # 提取全部文本
    for i, page_layout in enumerate(extract_pages(filename)):
        # 如果指定了页码范围，跳过范围外的页
        if page_numbers is not None and i not in page_numbers:
            continue
        for element in page_layout:
            if isinstance(element, LTTextContainer):
                full_text += element.get_text() + '\n'
    # 按空行分隔，将文本重新组织成段落
    lines = full_text.split('\n')
    for text in lines:
        if len(text) >= min_line_length:
            buffer += (' ' + text) if not text.endswith('-') else text.strip('-')
        elif buffer:
            paragraphs.append(buffer)
            buffer = ''
        print(text)
    if buffer:
        paragraphs.append(buffer)
    return paragraphs


def process_pdf_and_insert_to_milvus(pdf_path):
    """
    处理PDF文件，读取内容，分割段落，并将段落插入Milvus。
    
    参数:
    pdf_path (str): 要处理的PDF文件路径。
    """
    # 读取PDF并分割段落
    paragraphs = extract_text_from_pdf(pdf_path)

    # 将段落列表插入Milvus
    insert_texts_to_milvus(paragraphs)
    print("All paragraphs from the PDF have been inserted into Milvus.")


if __name__ == "__main__":
    pdf_file_path = "/Users/admin/Downloads/www_milvus-io_com__.pdf"
    process_pdf_and_insert_to_milvus(pdf_file_path)
